import torch
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.figure import Figure

# import torch.distributions  # 模块提供了许多概率分布的实现，可用于生成随机样本和计算概率
"""
    multinomial 类是一个多项分布（Multinomial Distribution）的实现。多项分布是一种离散概率分布
    它描述了在 n 次独立试验中，每次试验有 k 个可能的结果，每个结果有一个对应的概率，并且试验结果的总和为 n
"""
from torch.distributions import multinomial

# 每个可能结果的概率向量(6个可能的结果)，每个结果的概率为   1 / 6
# 结果为对应的6元素的向量，向量的每个数据对应相应的结果的次数
fair_probs = torch.ones([6]) / 6

# 根据给定的结果的概率向量抽取一个样本，从同一分布中生成多个样本
X = multinomial.Multinomial(1, fair_probs).sample()

print(X.shape)  # torch.Size([6])

print(X)  # tensor([0., 0., 1., 0., 0., 0.])
Y = multinomial.Multinomial(10, fair_probs).sample()

print(Y)  # tensor([2., 1., 2., 2., 1., 2.])

# 1000个采样
counts = multinomial.Multinomial(1000, fair_probs).sample()
counts /= 1000  # 占比即概率
print(counts)  # tensor([0.1690, 0.1720, 0.1630, 0.1710, 0.1480, 0.1770])

counts1 = multinomial.Multinomial(10, fair_probs).sample((500,))  # (500,6) 500组，每组10次的抽样结果
cum_counts = counts1.cumsum(dim=0)  # 取出第一维数据，并求和
print('cum ', cum_counts.shape)  # torch.Size([500, 6])
col_sum = cum_counts.sum(dim=1, keepdim=True)  # 每个第一维度中的，第二维度数据相加
col_sum1 = cum_counts.sum(dim=1)  # 每个第一维度中的，第二维度数据相加
print('cum ', col_sum.shape)
print('cum1 ', col_sum1.shape)
estimates = cum_counts / col_sum  # 求占比/概率

fig, axis = plt.subplots(nrows=1, ncols=1)
axis: Axes = axis  # 类型注解，加上这行代码后便可以提示
fig: Figure = fig

print(estimates[:0])  # 取从起始索引到0索引，不包含0，所以为空

for i in range(6):
    # 取出各个点数随着次数增加的概率值
    axis.plot(estimates[:, i:i + 1].numpy().flatten(), label=str(i + 1))

axis.axhline(y=0.167, color='black', linestyle='dashed')  # 绘制一个水平参考线
axis.legend()  # 显示图例-->对符号的说明
plt.show()
